/*
 *   Copyright 2021 zzh
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *
 */

package com.aduib.boot.common.buffer;

import java.util.Collections;
import java.util.LinkedList;
import java.util.List;

/**
 * @description: 大数据的写入
 * @author: zzh
 * @date: 2021/12/1 14:46
 */
public final class BigBuffer extends AbstractReferenceCounted {
  private static byte[] EMPTY = new byte[0];
  private final BufferPool bufferPool;
  private Buffer curWritingBlock;
  private List<Buffer> writeBlockLst = Collections.emptyList();

  public BigBuffer(BufferPool bufferPool) {
    super();
    this.bufferPool = bufferPool;
    curWritingBlock = bufferPool.allocate(bufferPool.getChunkSize());
  }

  public BigBuffer(BufferPool bufferPool, int initialCapacity) {
    super();
    this.bufferPool = bufferPool;
    curWritingBlock = bufferPool.allocate(initialCapacity);
  }

  public Buffer checkWriteBuffer(int capacity) {
    if (capacity > curWritingBlock.remaining()) {
      addBlock(curWritingBlock);
      curWritingBlock = bufferPool.allocate(capacity);
      return curWritingBlock;
    } else {
      return curWritingBlock;
    }
  }

  public int getBlockCount() {
    return writeBlockLst.size() + 1;
  }

  private void addBlock(Buffer buffer) {
    if (writeBlockLst.isEmpty()) {
      writeBlockLst = new LinkedList<Buffer>();
    }
    writeBlockLst.add(buffer);
  }

  public Buffer getCurWritingBlock() {
    return curWritingBlock;
  }

  public List<Buffer> getWriteBlockLst() {
    return writeBlockLst;
  }

  public void clear() {
    curWritingBlock = null;
    writeBlockLst.clear();
    writeBlockLst = null;
  }

  public Buffer write(byte[] src) {
    int offset = 0;
    int remains = src.length;
    while (remains > 0) {
      int writeable = curWritingBlock.remaining();
      if (writeable >= remains) {
        // can write whole srce
        curWritingBlock.put(src, offset, remains);
        break;
      } else {
        // can write partly
        curWritingBlock.put(src, offset, writeable);
        offset += writeable;
        remains -= writeable;
        addBlock(curWritingBlock);
        curWritingBlock = bufferPool.allocate(bufferPool.getChunkSize());
        continue;
      }
    }
    return curWritingBlock;
  }

  public byte[] writeToByteArrayAndRecycle() {
    BigBuffer bigBuffer = this;
    try {

      int size = 0;
      List<Buffer> blockes = bigBuffer.getWriteBlockLst();
      if (!bigBuffer.getWriteBlockLst().isEmpty()) {
        for (Buffer curBuf : blockes) {
          curBuf.switchReadMode();
          size += curBuf.remaining();
        }
      }
      Buffer curBuf = bigBuffer.getCurWritingBlock();
      curBuf.switchReadMode();
      if (curBuf.hasRemaining()) {
        size += curBuf.remaining();
      }
      if (size > 0) {
        int offset = 0;
        byte[] all = new byte[size];
        if (!bigBuffer.getWriteBlockLst().isEmpty()) {
          for (Buffer tBuf : blockes) {

            BufferUtils.copy(tBuf, 0, all, offset, tBuf.remaining());
            offset += tBuf.remaining();

            bufferPool.recycle(tBuf);
          }
        }
        Buffer tBuf = bigBuffer.getCurWritingBlock();
        if (tBuf.hasRemaining()) {
          BufferUtils.copy(tBuf, 0, all, offset, tBuf.remaining());
          bufferPool.recycle(tBuf);
        }
        return all;
      }

    } finally {
      bigBuffer.clear();
    }

    return EMPTY;
  }

  @Override
  protected void deallocate() {
    clear();
  }
}
